/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.core;

import com.lvyh.lightframe.tracer.core.util.TraceIdGenerator;
import com.lvyh.lightframe.tracer.core.samplers.Sampler;
import com.lvyh.lightframe.tracer.core.util.StringUtils;
import io.opentracing.SpanContext;
import lombok.Data;

import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

@Data
public class SimpleSpanContext implements SpanContext {

    private static final Charset CHARSET = Charset.forName("UTF-8");

    private static final String RPC_ID_SEPARATOR = ".";

    private static final byte NO_ENTRY = 0, ENTRY = 1;

    private String traceId = "";

    private String spanId = "";

    private String parentId = "";

    private boolean isSampled = false;

    private static final String TRACE_ID_ALIAS = "tcid";

    private static final String SPAN_ID_ALIAS = "spid";

    private static final String PARENT_SPAN_ID_ALIAS = "pspid";

    private static final String SAMPLE_ALIAS = "sample";

    /**
     * The serialization system transparently passes the prefix of the attribute key
     */
    private static final String SYS_BAGGAGE_PREFIX_KEY = "_sys_";

    private Sampler sampler;

    private Map<String, String> baggage = new ConcurrentHashMap<String, String>();

    private AtomicInteger childContextIndex = new AtomicInteger(0);

    public SimpleSpanContext() {
        this(StringUtils.EMPTY_STRING, StringUtils.EMPTY_STRING, null, false);
    }

    public SimpleSpanContext(String traceId, String spanId, String parentId, boolean isSampled) {
        this.traceId = traceId;
        this.spanId = spanId;
        this.parentId = StringUtils.isBlank(parentId) ? this.genParentSpanId(spanId) : parentId;
        this.isSampled = isSampled;
    }

    public SimpleSpanContext(String traceId, String spanId) {
        this(traceId, spanId, null, false);
    }

    public SimpleSpanContext(String traceId, String spanId, String parentId) {
        this(traceId, spanId, parentId, false);
    }

    @Override
    public Iterable<Map.Entry<String, String>> baggageItems() {
        Map<String, String> allBaggage = new HashMap<String, String>();
        if (this.baggage != null && this.baggage.size() > 0) {
            allBaggage.putAll(this.baggage);
        }
        return allBaggage.entrySet();
    }

    private String genParentSpanId(String spanId) {
        return (StringUtils.isBlank(spanId) || spanId.lastIndexOf(RPC_ID_SEPARATOR) < 0) ? StringUtils.EMPTY_STRING
                : spanId.substring(0, spanId.lastIndexOf(RPC_ID_SEPARATOR));
    }

    public SimpleSpanContext addBaggage(Map<String, String> sysBaggage) {
        if (sysBaggage != null && sysBaggage.size() > 0) {
            this.baggage.putAll(sysBaggage);
        }
        return this;
    }

    public static SimpleSpanContext rootStart() {
        return rootStart(false);
    }

    public static SimpleSpanContext rootStart(boolean isSampled) {
        String traceId = TraceIdGenerator.generate();
        return new SimpleSpanContext(traceId, "0", StringUtils.EMPTY_STRING, isSampled);
    }


    public String serialize() {
        StringBuilder stringBuilder = new StringBuilder();
        stringBuilder.append(TRACE_ID_ALIAS).append(StringUtils.EQUAL).append(traceId);
        stringBuilder.append(StringUtils.AND);
        stringBuilder.append(SPAN_ID_ALIAS).append(StringUtils.EQUAL).append(spanId);
        stringBuilder.append(StringUtils.AND);
        stringBuilder.append(PARENT_SPAN_ID_ALIAS).append(StringUtils.EQUAL).append(parentId);
        stringBuilder.append(StringUtils.AND);
        stringBuilder.append(SAMPLE_ALIAS).append(StringUtils.EQUAL).append(isSampled);
        stringBuilder.append(StringUtils.AND);
        if (this.baggage.size() > 0) {
            stringBuilder.append(StringUtils.mapToString(this.baggage));
        }
        return stringBuilder.toString();
    }


    public static SimpleSpanContext deserialize(String deserializeValue) {
        if (StringUtils.isBlank(deserializeValue)) {
            return SimpleSpanContext.rootStart();
        }

        String traceId = TraceIdGenerator.generate();
        String spanId = SimpleTracer.ROOT_SPAN_ID;
        String parentId = StringUtils.EMPTY_STRING;
        boolean sampled = false;

        Map<String, String> baggage = new HashMap<String, String>();

        Map<String, String> spanContext = new HashMap<String, String>();
        StringUtils.stringToMap(deserializeValue, spanContext);

        for (Map.Entry<String, String> entry : spanContext.entrySet()) {
            String key = entry.getKey();
            String value = entry.getValue();
            if (StringUtils.isBlank(key) || StringUtils.isBlank(value)) {
                continue;
            }
            if (TRACE_ID_ALIAS.equals(key)) {
                traceId = value;
                continue;
            }
            if (SPAN_ID_ALIAS.equals(key)) {
                spanId = value;
                continue;
            }
            if (PARENT_SPAN_ID_ALIAS.equals(key)) {
                parentId = value;
                continue;
            }
            if (SAMPLE_ALIAS.equals(key)) {
                sampled = Boolean.parseBoolean(value);
                continue;
            }
            int sysIndex = key.indexOf(SYS_BAGGAGE_PREFIX_KEY);
            if (sysIndex == 0) {
                //must have a prefix
                String sysKey = key.substring(SYS_BAGGAGE_PREFIX_KEY.length());
                baggage.put(sysKey, value);
            } else {
                baggage.put(key, value);
            }
        }
        SimpleSpanContext simpleSpanContext = new SimpleSpanContext(traceId, spanId, parentId, sampled);
        if (baggage.size() > 0) {
            simpleSpanContext.addBaggage(baggage);
        }
        return simpleSpanContext;
    }

    /**
     * Get the ID of the next sub context
     *
     * @return next spanId
     */
    public String nextChildContextId() {
        return this.spanId + RPC_ID_SEPARATOR + childContextIndex.incrementAndGet();
    }

    @Override
    public String toSpanId() {
        return this.getSpanId();
    }

    @Override
    public String toTraceId() {
        return this.getTraceId();
    }
}
